import os
import sys
import json
import numpy as np
import tensorrt as trt
import ctypes
from tensorrt import plugin as trt_plugin

# === 路径加入工程根目录 ===
PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
if PROJECT_ROOT not in sys.path:
    sys.path.insert(0, PROJECT_ROOT)

ONNX_PATH        = os.path.join(PROJECT_ROOT, 'onnx/yolov11n_int8_ALLConcat.onnx')
PPQ_JSON_PATH    = os.path.join(PROJECT_ROOT, 'onnx/yolov11n_int8_ALLConcat.json')
ENGINE_OUTPUT_PATH = os.path.join(PROJECT_ROOT, 'engine/yolov11n_int8_ALLConcat.engine')

TRT_LOGGER = trt.Logger()
EXPLICIT_BATCH = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)

def GiB(val):
    return val * 1 << 30

def json_load(filename):
    with open(filename) as json_file:
        data = json.load(json_file)
    return data

def setDynamicRange(network, json_file):
    quant_param_json = json_load(json_file)
    act_quant = quant_param_json.get("act_quant_info", {})
    for i in range(network.num_inputs):
        input_tensor = network.get_input(i)
        if input_tensor.name in act_quant:
            value = act_quant[input_tensor.name]
            tensor_max = abs(value)
            tensor_min = -abs(value)
            input_tensor.set_dynamic_range(tensor_min, tensor_max)
    for i in range(network.num_layers):
        layer = network.get_layer(i)
        for output_index in range(layer.num_outputs):
            tensor = layer.get_output(output_index)
            if tensor is not None and tensor.name in act_quant:
                value = act_quant[tensor.name]
                tensor_max = abs(value)
                tensor_min = -abs(value)
                tensor.set_dynamic_range(tensor_min, tensor_max)

def build_engine(onnx_file, json_file, engine_file):
    import numpy as np
    import ctypes
    import os

    try:
        ctypes.CDLL("libnvinfer_plugin.so")
    except OSError:
        pass

    builder = trt.Builder(TRT_LOGGER)
    network = builder.create_network(EXPLICIT_BATCH)
    parser = trt.OnnxParser(network, TRT_LOGGER)
    config = builder.create_builder_config()
    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, GiB(1))
    config.set_flag(trt.BuilderFlag.INT8)

    if not os.path.exists(onnx_file):
        quit(f'ONNX file {onnx_file} not found')
    with open(onnx_file, 'rb') as model:
        if not parser.parse(model.read()):
            print('ERROR: Failed to parse the ONNX file.')
            for error in range(parser.num_errors):
                print(parser.get_error(error))
            return None

    setDynamicRange(network, json_file)

    num_outputs = network.num_outputs
    print("Your ONNX network outputs:")
    for i in range(num_outputs):
        print(f"{i}: {network.get_output(i).name}, shape={network.get_output(i).shape}")

    output = network.get_output(0)
    print("Raw output shape:", output.shape)
    batch, ch, nbox = output.shape
    assert ch >= 5, "output第二维少于5，模型格式不符"
    num_classes = ch - 4

    # boxes: [1,4,8400] -> [1,8400,4]
    boxes_slice = network.add_slice(output, start=(0,0,0), shape=(batch,4,nbox), stride=(1,1,1)).get_output(0)
    boxes_shuffle = network.add_shuffle(boxes_slice)
    boxes_shuffle.second_transpose = (0,2,1)
    boxes_for_nms = boxes_shuffle.get_output(0)

    # scores: [1,80,8400] -> [1,8400,80]
    scores_slice = network.add_slice(output, start=(0,4,0), shape=(batch,num_classes,nbox), stride=(1,1,1)).get_output(0)
    scores_shuffle = network.add_shuffle(scores_slice)
    scores_shuffle.second_transpose = (0,2,1)
    scores_for_nms = scores_shuffle.get_output(0)

    network.unmark_output(output)

    print("boxes shape (for NMS):", boxes_for_nms.shape)
    print("scores shape (for NMS):", scores_for_nms.shape)

    # -------- 插入 EfficientNMS_TRT 插件 --------
    plugin_creator = None
    for c in trt.get_plugin_registry().plugin_creator_list:
        if c.name == "EfficientNMS_TRT":
            plugin_creator = c
            break
    assert plugin_creator is not None, "EfficientNMS_TRT plugin not found!"

    nms_plugin_fields = [
        trt.PluginField("background_class", np.array([-1], dtype=np.int32), trt.PluginFieldType.INT32),
        trt.PluginField("max_output_boxes", np.array([100], dtype=np.int32), trt.PluginFieldType.INT32),
        trt.PluginField("score_threshold", np.array([0.25], dtype=np.float32), trt.PluginFieldType.FLOAT32),
        trt.PluginField("iou_threshold", np.array([0.45], dtype=np.float32), trt.PluginFieldType.FLOAT32),
        trt.PluginField("score_activation", np.array([0], dtype=np.int32), trt.PluginFieldType.INT32), # 默认不sigmoid
        trt.PluginField("box_coding", np.array([0], dtype=np.int32), trt.PluginFieldType.INT32),       # 0: [x1,y1,x2,y2]
    ]
    nms_plugin_fields = trt.PluginFieldCollection(nms_plugin_fields)
    nms_plugin = plugin_creator.create_plugin(name="EfficientNMS_TRT", field_collection=nms_plugin_fields)

    nms_layer = network.add_plugin_v2([boxes_for_nms, scores_for_nms], nms_plugin)

    nms_output_names = ["num_dets", "dets_boxes", "dets_scores", "dets_classes"]
    for i in range(nms_layer.num_outputs):
        nms_layer.get_output(i).name = nms_output_names[i]
        network.mark_output(nms_layer.get_output(i))

    # ----------- 构建TensorRT引擎 -----------
    serialized_engine = builder.build_serialized_network(network, config)
    assert serialized_engine is not None, 'Network build failed!'

    runtime = trt.Runtime(TRT_LOGGER)
    engine = runtime.deserialize_cuda_engine(serialized_engine)
    with open(engine_file, "wb") as f:
        f.write(engine.serialize())
    print(f"\033[1;32mEngine regenerated with NMS: {engine_file}\033[0m")
    print("输出节点:", [network.get_output(i).name for i in range(network.num_outputs)])
    print("=== Final outputs from engine:")
    for i in range(network.num_outputs):
        out = network.get_output(i)
        print(f"{i}: name={out.name}, shape={out.shape}, dtype={out.dtype}")
    
    
    
if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser(description='Build TRT engine with NMS from ONNX. For 1*84*8400 output.')
    parser.add_argument('--onnx', type=str,default=ONNX_PATH ,help='ONNX file')
    parser.add_argument('--qparam_json', type=str,default=PPQ_JSON_PATH, help='quantization config json')
    parser.add_argument('--engine', type=str, default=ENGINE_OUTPUT_PATH, help='Output TRT engine file')
    args = parser.parse_args()

    build_engine(args.onnx, args.qparam_json, args.engine)